// Copyright (c) Microsoft Corporation. All rights reserved.

using System;
using System.IO;
using System.Reflection;
using System.Runtime.InteropServices;
using System.Runtime.Serialization;
using System.Text;
using System.Threading;
using Microsoft.Win32.SafeHandles;
using System.Diagnostics;

namespace Microsoft.Tools.CodeRunner
{
    // TODO: Security - AllowPartiallyTrusted Callers

    /// <summary>
	/// The exception that is thrown when one of the command line arguments provided 
	/// is not valid.
	/// </summary>
	[Serializable()]
	public sealed class RemoteEnvironmentException : ArgumentException
	{
		#region Public Instance Constructors

		/// <summary>
		/// Initializes a new instance of the <see cref="RemoteEnvironmentException" /> class.
		/// </summary>
		public RemoteEnvironmentException()
			: base()
		{
		}

		/// <summary>
		/// Initializes a new instance of the <see cref="RemoteEnvironmentException" /> class
		/// with a descriptive message.
		/// </summary>
		/// <param name="message">A descriptive message to include with the exception.</param>
		public RemoteEnvironmentException(string message)
			: base(message)
		{
		}

		/// <summary>
		/// Initializes a new instance of the <see cref="RemoteEnvironmentException" /> class
		/// with a descriptive message and an inner exception.
		/// </summary>
		/// <param name="message">A descriptive message to include with the exception.</param>
		/// <param name="innerException">A nested exception that is the cause of the current exception.</param>
		public RemoteEnvironmentException(string message, Exception innerException)
			: base(message, innerException)
		{
		}

		#endregion Public Instance Constructors

		#region Private Instance Constructors

		/// <summary>
		/// Initializes a new instance of the <see cref="RemoteEnvironmentException" /> class 
		/// with serialized data.
		/// </summary>
		/// <param name="info">The <see cref="SerializationInfo" /> that holds the serialized object data about the exception being thrown.</param>
		/// <param name="context">The <see cref="StreamingContext" /> that contains contextual information about the source or destination.</param>
		private RemoteEnvironmentException(SerializationInfo info, StreamingContext context)
			: base(info, context)
		{
		}

		#endregion Private Instance Constructors
	}

	/// <summary>
	/// Summary description for RemoteEnvironment.
	/// </summary>
	public class RemoteEnvironment : IDisposable
	{
		#region Private Data
		
		private ProcessInjector injector;
		private SafeFileHandle pipe;
		
		private const int maxMessageLength = 32768;
		private const int maxRetries = 1024;
		
		#endregion

        #region Public Constructors
        /// <summary>
        /// Initializes a new instances of the <see cref="RemoteEnvironment"/> class.  Note that if this method 
        /// fails there may be a temporary file left in the TEMP directory.
        /// </summary>
        /// <param name="processId"></param>
        public RemoteEnvironment(int processId)
        {
            // Create a pipe to communicate with the remote thread.  This must be done first so that the 
            // pipe is ready when the remote thread starts
            pipe = NativeMethods.CreateNamedPipe(
                @"\\.\pipe\coderunner$", PipeAccess.Duplex,
                PipeMode.TypeMessage | PipeMode.ReadModeMessage | PipeMode.Wait,
                255, maxMessageLength * Marshal.SizeOf(typeof(char)), maxMessageLength * Marshal.SizeOf(typeof(char)),
                0, IntPtr.Zero);

            if (pipe.IsInvalid)
            {
                throw new RemoteEnvironmentException("Unable to create pipe to remote process");
            }

            string dllFileName = Path.GetTempFileName();

            // Unpack the DLL to be injected from our resources and put in the temporary file
            using (BinaryReader br = new BinaryReader(Assembly.GetExecutingAssembly().GetManifestResourceStream(
                typeof(RemoteEnvironment), "CodeRunnerHelper.dll")))
            {
                using (BinaryWriter bw = new BinaryWriter(new FileStream(dllFileName, FileMode.Create)))
                {
                    bw.Write(br.ReadBytes((int)br.BaseStream.Length));
                }
            }

            // Inject the remote thread and pull in the DLL to create the other end of the pipe
            injector = new ProcessInjector(processId, dllFileName, "ControlRemoteEnvironment");

            // TODO: This will block until the injected thread connects. If anything goes wrong, such as the remote process dies, 
            // we'll hang. Let's not do that.
            bool connected = NativeMethods.ConnectNamedPipe(pipe, IntPtr.Zero) ?
                true : (Marshal.GetLastWin32Error() == Error.PipeConnected);
        }

        #endregion
		
		#region Native methods

		[Flags]
		private enum PipeMode
		{
			Wait = 0x00000000,
			NoWait = 0x00000001,
			ReadModeByte = 0x00000000,
			ReadModeMessage = 0x00000002,
			TypeByte = 0x00000000,
			TypeMessage = 0x00000004,
		}

		private enum PipeAccess
		{
			Inbound = 0x00000001,
			Outbound = 0x00000002,
			Duplex = 0x00000003
		}
		
		private class Error
		{
			public const int PipeConnected = 535;
		}

		private class NativeMethods
		{
			[DllImport("kernel32.dll", SetLastError = true)]
			public static extern bool SetNamedPipeHandleState(
				SafeFileHandle hPipe, ref PipeMode mode,
				IntPtr maxCollectionCount, IntPtr collectDataTimeout);

            [DllImport("kernel32.dll", SetLastError = true)]
            public static extern bool WriteFile(
                SafeFileHandle hFile, byte[] buffer, int numberOfBytesToWrite,
                out int numberOfBytesWritten, IntPtr overlapped);

            [DllImport("kernel32.dll", SetLastError = true)]
            public static extern bool ReadFile(
                SafeFileHandle hFile, byte[] buffer, int numberOfBytesToRead,
                out int numberOfBytesRead, IntPtr overlapped);

            [DllImport("kernel32.dll", SetLastError = true)]
            public static extern bool PeekNamedPipe(
                SafeFileHandle hFile, byte[] buffer, int bufferSize,
                out int bytesRead, out int bytesAvailable, out int bytesLeftThisMessage);

            [DllImport("kernel32.dll", SetLastError = true)]
			public static extern SafeFileHandle CreateNamedPipe(
				string name, PipeAccess access, PipeMode mode, 
				int maxInstances, int outBufferSize, int inBufferSize,
				int defaultTimeout, IntPtr securityDescriptor);

			[DllImport("kernel32.dll", SetLastError = true)]
			public static extern bool ConnectNamedPipe(SafeFileHandle hPipe, IntPtr overlapped);
		}

		#endregion
	
		#region IDisposable and Related Members

		/// <summary>
		/// Finalizer
		/// </summary>
		~RemoteEnvironment()
		{
			Dispose(false);
		}
		
		/// <summary>
		/// Dispose method
		/// </summary>
		public void Dispose()
		{
			Dispose(true);
			GC.SuppressFinalize(this);
		}

		private void Dispose(bool disposing)
		{
			// If we are finalizing it's not safe to access any object members; this includes if an exception 
			// was thrown in the constructor.
			// If we are explicitly disposing then we assume that everything was initialized correctly if injector
			// is not null as that was the last thing created in the constructor.
			if (disposing && injector != null)
			{
				// First, we need to cause the DLL in the remote process to unload by sending a 
				// terminate message to the remote thread.
				byte[] bytes = Encoding.Unicode.GetBytes("terminate\0");
				int bytesWritten;
				
				NativeMethods.WriteFile(pipe, bytes, bytes.Length, out bytesWritten, IntPtr.Zero);
							
				// Now we are done with the pipe
				pipe.Close();
				pipe = null;

				// Wait forever for the remote thread to exit.  It's vital that it has exited before we release
				// the ProcessInjector object, or else the remote thread will return into freed memory and we 
				// bring down the remote process!  
				injector.WaitForExit();

				try
				{
					File.Delete(injector.DllFileName);
				}
				catch (SystemException)
				{
					// This is just a precaution, because nothing throws exceptions like I/O operations!
				}

				injector.Dispose();
				injector = null;
			}
		}
		#endregion	

        #region Private Methods
        private void SendMessage(string message)
        {
            byte[] bytes = Encoding.Unicode.GetBytes(message);
            int bytesWritten;

            if (!NativeMethods.WriteFile(pipe, bytes, bytes.Length, out bytesWritten, IntPtr.Zero) || 
                bytesWritten != bytes.Length)
                throw new RemoteEnvironmentException();
        }

        private string ReceiveMessage()
        {
            byte[] bytes = new byte[maxMessageLength];
            int bytesRead;
            int bytesAvailable;
            int bytesBytesLeftThisMessage;

            Thread.Sleep(0);

            for (int i = maxRetries; i > 0; i--)
            {
                if (NativeMethods.PeekNamedPipe(pipe, null, 0, out bytesRead, out bytesAvailable, out bytesBytesLeftThisMessage) &&
                    bytesAvailable > 0)
                {
                    int numberOfBytesRead;

                    NativeMethods.ReadFile(pipe, bytes, bytesAvailable, out numberOfBytesRead, IntPtr.Zero);

                    Array.Resize(ref bytes, numberOfBytesRead - sizeof(char));

                    return Encoding.Unicode.GetString(bytes);
                }

                Thread.Sleep(10);
            }

            //throw new TimeoutException();
            return null;
        }

        #endregion

		#region Public Methods
        /// <summary>
        /// Set an environment variable in remote process
        /// </summary>
        /// <param name="name">Name of the environment variable</param>
        /// <param name="val">Value to set the environment variable to</param>
        public void SetEnvironmentVariable(string name, string val)
        {
            if (pipe == null)
                throw new InvalidOperationException("Object has been disposed");

            // Check maximum length of environment variable and name. The value includes space for the "set" command.
            if (name.Length + val.Length >= 32760)
                throw new ArgumentException("Length of name and value must be < 32760");

            SendMessage(String.Format("setvar\r\n{0}={1}\0", name, val));
            ReceiveMessage();
        }

        /// <summary>
        /// Get an environment variable in remote process
        /// </summary>
        /// <param name="name">Name of the environment variable</param>
        public string GetEnvironmentVariable(string name)
        {
            if (pipe == null)
                throw new InvalidOperationException("Object has been disposed");

            SendMessage(String.Format("getvar\r\n{0}\0", name));
            return ReceiveMessage();
        }

        /// <summary>
        /// Get an array of the environment variable names
        /// </summary>
        /// <param name="names">Array of names of the environment variables</param>
        public string[] GetEnvironmentVariableNames()
        {
            if (pipe == null)
                throw new InvalidOperationException("Object has been disposed");

            SendMessage(String.Format("getvarnames\0"));
            
            string result = ReceiveMessage();

            return result.Split(new string[] { "\r\n" }, StringSplitOptions.RemoveEmptyEntries);
        }

        /// <summary>
        /// Get or set the current directory of a remote process
        /// </summary>
        public string CurrentDirectory
        {
            get
            {
                if (pipe == null)
                    throw new InvalidOperationException("Object has been disposed");

                SendMessage(String.Format("getcd\0"));
                return ReceiveMessage();
            }
            set
            {
                if (pipe == null)
                    throw new InvalidOperationException("Object has been disposed");

                SendMessage(String.Format("setcd\r\n{0}\0", value));
                ReceiveMessage();
            }
        }
        #endregion
	}
}
